"""
    时间：2019/10/24
    作者：徐飞
    版本：01
    内容：20期第三次作业
        第三题
"""


import cv2 as cv
import copy


cap = cv.VideoCapture(r"vtest.avi")
size = (int(cap.get(cv.CAP_PROP_FRAME_WIDTH) * 0.8), int(cap.get(cv.CAP_PROP_FRAME_HEIGHT) * 0.8))
element = cv.getStructuringElement(cv.MORPH_CROSS, (3, 3))
# 角点检测参数
feature_param = dict(maxCorners=100, qualityLevel=0.3, minDistance=0.7, blockSize=7)
# l-k光流法参数
LK_param = dict(winSize=(15, 15), maxLevel=2,
                criteria=(cv.TERM_CRITERIA_COUNT | cv.TERM_CRITERIA_EPS, 10, 0.03))
# 第一帧特征点
ret, prev = cap.read()
prevGray = cv.cvtColor(prev, cv.COLOR_BGR2GRAY)
p0 = cv.goodFeaturesToTrack(prevGray, mask=None, **feature_param)

# 循环进行下一帧
while True:
    ret, frame = cap.read()
    if not ret:    # 没读到当前帧
        break
    gray = cv.cvtColor(frame, cv.COLOR_BGR2GRAY)
    # 计算流光
    p1, st, err = cv.calcOpticalFlowPyrLK(prevGray, gray, p0, None, **LK_param)
    # 选择好的跟踪点
    goodPoints = p1[st == 1]
    goodPrevPoint = p0[st == 1]
    # 在结果图像中迭加出特征点和光流向量
    res = copy.deepcopy(frame)
    drawColor = (0, 0, 255)
    drawColor1 = (255, 0, 0)
    for i, (cur, prev) in enumerate(zip(goodPoints, goodPrevPoint)):
        x0, y0 = cur.ravel()
        x1, y1 = prev.ravel()
        cv.line(res, (x0, y0), (x1, y1), drawColor1)
        cv.circle(res, (x0, y0), 3, drawColor)
    # 更新上一帧
    p0 = goodPoints.reshape(-1, 1, 2)
    # 显示结果图片
    cv.imshow("检测结果", res)
    if cv.waitKey(100) & 0xff == 27:   # 每30ms一帧
        break
cap.release()
cv.destroyAllWindows()
